import user  # look for .pythonrc.py for user init
import math
import ROOT
import PyCintex
import AthenaROOTAccess.transientTree

from PyAnalysisUtils import PDG
CollectionTree = ROOT.AthenaROOTAccess.TChainROOTAccess('CollectionTree')
#'''My Private Production'''
#CollectionTree.Add('/tmp/narayan/mc11_7TeV.109300.AlpgenJimmyZeebbNp0_nofilter.evgen.EVNT.e835_tid619473_00/*root*')
#CollectionTree.Add('/tmp/narayan/mc12_7TeV.109380.AlpgenJimmyZeebbNp0_3Leptons.evgen.EVNT.e835_tid454689_00/*root*')
CollectionTree.Add('/tmp/narayan/mc11_7TeV.109300.AlpgenJimmyZeebbNp0_nofilter.merge.AOD.e835_s1310_s1300_r3043_r2993_tid645621_00/*root*')
#CollectionTree.Add('/tmp/narayan/alpgenjimmy.pool.root')
#CollectionTree.Add('/tmp/narayan/mc11_7TeV.105200.T1_McAtNlo_Jimmy.merge.AOD.e835_s1310_s1300_r3257_r3063_tid729124_00/AOD.729124._000020.pool.root.1')
#CollectionTree.Add('/tmp/narayan/mc11_7TeV.126416.AlpgenJimmyZeeccNp2_nofilter.merge.AOD.e1008_s1310_s1300_r3043_r2993_tid679751_00/*root*')
#CollectionTree.Add('/tmp/narayan/mc11_7TeV.126414.AlpgenJimmyZeeccNp0_nofilter.evgen.EVNT.e1008_tid636322_00/*root*')

tt = AthenaROOTAccess.transientTree.makeTree(CollectionTree)

def toiter (beg, end):
       while beg != end:
           yield beg.__deref__()
           beg.__preinc__()
       return
def DeltaR(p,q):
    from math import sqrt
    return sqrt((p.momentum().eta() - q.momentum().eta())**2 + (p.momentum().phi() - q.momentum().phi())**2)

def IsBHadron(p):
    status = p.status()
    mpdg = abs(p.pdg_id())
    if( ( 500 < mpdg and mpdg < 599 )    or \
            ( 10500 < mpdg and mpdg < 10599 )  or \
            (  5000 < mpdg and mpdg < 5999  )  or \
            ( 20500 < mpdg and mpdg < 20599 ) ):
            return True
    else:
        return False
    
    

def IsCHadron(p):
    status = p.status()
    mpdg = abs(p.pdg_id())
    #if(status == 196 or status == 197 or status ==198): 
    if( ( 400  < mpdg and mpdg < 499 )   or \
            ( 10400 < mpdg and mpdg < 10499 ) or \
            (  4000 < mpdg and mpdg < 4999  ) or \
            ( 20400 < mpdg and mpdg < 20499 ) ):
        return True
    else:
        return False

def GetParents(p):
    parentList= []
    if(p.production_vertex()):
        for parent in toiter(p.production_vertex().particles_begin(0),p.production_vertex().particles_end(0)):
            if(abs(parent.pdg_id()) == abs(p.pdg_id())):
                return GetParents(parent)
            else:
                parentList +=[parent]
    return parentList

def GetChildren(p):
    childList= []
    if(p.end_vertex()):
        for child in toiter(p.end_vertex().particles_begin(1),p.end_vertex().particles_end(1)):
            if (child.pdg_id() == p.pdg_id()):
                return GetChildren(child)
            else:
                childList +=[child]
    return childList


def IsFinalState(p,type):
    if(p.end_vertex()):
        for child in toiter(p.end_vertex().particles_begin(1),p.end_vertex().particles_end(1)):
            if(child.barcode()>100000):
                continue
            pdgid   = child.pdg_id()
            rest1   = abs(pdgid) % 1000
            rest2   = abs(pdgid) % 10000

            if(type == 5):
                if( (rest2>=5000 and rest2 <6000) or (rest1 >=500 and rest1 < 600)):
                    return False
                if(child.barcode() < 100000):
                    if( not IsFinalState(child,type)):
                        return False
            if(type == 4):
                if( (rest2>=4000 and rest2 < 5000) or (rest1>=400 and rest1 <500)):
                    return False
                if(child.barcode() < 100000):
                    if( not IsFinalState(child,type)):
                        return False
    return True
    

nBHadrons = 0.0
nCHadrons = 0.0

nSemileptB = 0.0
nSemileptC = 0.0

h_bptCorrel = ROOT.TH2F("bptcorrel","b quark pt correlation;b1;b1",500,0,500,500,0,500)
h_betaCorrel= ROOT.TH2F("bEtacorrel","b quark eta correlation; b1;b2",100,-5,5,100,-5,5)    
h_zvec      = ROOT.TH3F("Z3Vector",";pt;eta;phi; #eta",500,0,500,100,-5,5,100,-5,5)

for iEvt in range(tt.GetEntriesFast()):
    tt.GetEntry(iEvt)
    mcc = tt.SpclMC
    ev = mcc.genEvent()
    #ev  = tt.GEN_EVENT[0]
    #ElColl = tt.ElectronAODCollection
    print "Event==============================================================: ",iEvt 

    m1 = 0
    m2 = 0

    bpt     = 0
    beta    = 0
    bbarpt  = 0
    bbareta = 0

    particleList = []

    for p in toiter(ev.particles_begin(),ev.particles_end()):
        if(abs(p.pdg_id()) ==23):
            print p.pdg_id(),p.status()
